{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "d=64\n",
    "nb = 100000\n",
    "nq = 10000\n",
    "np.random.seed(1234)\n",
    "xb = np.random.random((nb, d)).astype('float32')\n",
    "xb[:, 0] += np.arange(nb) / 1000.\n",
    "xq = np.random.random((nq, d)).astype('float32')\n",
    "xq[:, 0] += np.arange(nq) / 1000."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100000\n"
     ]
    }
   ],
   "source": [
    "import faiss\n",
    "#####brute search\n",
    "index = faiss.IndexFlatL2(d)\n",
    "#print(index.is_trained)\n",
    "index.add(xb)\n",
    "print(index.ntotal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[  0 393 363  78]\n",
      " [  1 555 277 364]\n",
      " [  2 304 101  13]\n",
      " [  3 173  18 182]\n",
      " [  4 288 370 531]]\n",
      "[[0.        7.1751733 7.207629  7.2511625]\n",
      " [0.        6.3235645 6.684581  6.7999454]\n",
      " [0.        5.7964087 6.391736  7.2815123]\n",
      " [0.        7.2779055 7.5279865 7.6628466]\n",
      " [0.        6.7638035 7.2951202 7.3688145]]\n"
     ]
    }
   ],
   "source": [
    "k = 4\n",
    "D, I = index.search(xb[:5], k)\n",
    "#print(xb[:5])\n",
    "print(I)\n",
    "print(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 381  207  210  477]\n",
      " [ 526  911  142   72]\n",
      " [ 838  527 1290  425]\n",
      " [ 196  184  164  359]\n",
      " [ 526  377  120  425]]\n",
      "[[ 9900 10500  9309  9831]\n",
      " [11055 10895 10812 11321]\n",
      " [11353 11103 10164  9787]\n",
      " [10571 10664 10632  9638]\n",
      " [ 9628  9554 10036  9582]]\n"
     ]
    }
   ],
   "source": [
    "D, I = index.search(xq, k)\n",
    "print(I[:5])\n",
    "print(I[-5:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#####inverted file\n",
    "nlist=100\n",
    "quantizer = faiss.IndexFlatL2(d)\n",
    "index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)\n",
    "####METRIC_L2 is L2 distance, IP means inner product\n",
    "assert not index.is_trained\n",
    "index.train(xb)\n",
    "assert index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 9900  9309  9810 10048]\n",
      " [11055 10895 10812 11321]\n",
      " [11353 10164  9787 10719]\n",
      " [10571 10664 10632 10203]\n",
      " [ 9628  9554  9582 10304]]\n"
     ]
    }
   ],
   "source": [
    "###nprobe=1\n",
    "index.add(xb)\n",
    "D, I = index.search(xq, k)\n",
    "print(I[-5:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 9900 10500  9309  9831]\n",
      " [11055 10895 10812 11321]\n",
      " [11353 11103 10164  9787]\n",
      " [10571 10664 10632  9638]\n",
      " [ 9628  9554 10036  9582]]\n"
     ]
    }
   ],
   "source": [
    "index.nprobe = 10\n",
    "D, I = index.search(xq, k)\n",
    "print(I[-5:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[   0   78  608  159]\n",
      " [   1 1063  555  380]\n",
      " [   2  304  134   46]\n",
      " [   3   64  773  265]\n",
      " [   4  288  827  531]]\n",
      "[[1.6157436 6.1152253 6.4348025 6.564184 ]\n",
      " [1.389575  5.6771317 5.9956017 6.486294 ]\n",
      " [1.7025063 6.121688  6.189084  6.489888 ]\n",
      " [1.8057687 6.5440307 6.6684756 6.859398 ]\n",
      " [1.4920276 5.79976   6.190908  6.3791513]]\n"
     ]
    }
   ],
   "source": [
    "###pq to decrease the momery\n",
    "nlist = 100\n",
    "m=8\n",
    "k=4\n",
    "quantizer=faiss.IndexFlatL2(d)\n",
    "index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)\n",
    "#the last 8 specifies that each sub-vector is encoded as 8 bits\n",
    "index.train(xb)\n",
    "index.add(xb)\n",
    "D, I = index.search(xb[:5], k)\n",
    "print(I)\n",
    "print(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 9900  8746  9853 10437]\n",
      " [10494 10507 11373  9014]\n",
      " [10719 11291 10424 10138]\n",
      " [10122  9638 11113 10630]\n",
      " [ 9229 10304  9644 10370]]\n"
     ]
    }
   ],
   "source": [
    "index.nprobe = 10\n",
    "D, I = index.search(xq, k)\n",
    "print(I[-5:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18446744073709551616"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2**64"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
